//
// Created by syshen on 2021/7/7.
//

#ifndef OCR_BASE_CRNN_H
#define OCR_BASE_CRNN_H

#include <map>
#include <fstream>
#include <sstream>
#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <functional>
#include <memory>
#include <mutex>
#include <string>
#include <unordered_map>
#include <vector>

using namespace std;

class BaseCrnn {
public:
    BaseCrnn(int device_id, bool trt_serialize):device_id(device_id), trt_serialize(trt_serialize) {};
    BaseCrnn(BaseCrnn& /*other*/) noexcept = 0; // copy constructors
    BaseCrnn(BaseCrnn&& /*other*/) noexcept = 0; // move constructors
    BaseCrnn& operator=(BaseCrnn& ) noexcept = 0;
    BaseCrnn& operator=(BaseCrnn&& ) noexcept = 0;

    virtual int init(const std::string &model_file) = 0;

    virtual std::vector <std::string> extract(const cv::Mat &inputs) = 0;

    virtual ~BaseCrnn() {};

public:
    virtual bool getchar(const std::string &key_file) {
        std::ifstream fin(key_file);   // filename: xly2016I.txt  "../src/utils/ocr_keys.txt"
        std::stringstream buffer;            // stringstream object
        if (fin) // 有该文件
        {
            std::string line;
            while (getline(fin, line)) // line中不包括每行的换行符
            {
                line.erase(std::remove(line.begin(), line.end(), '\n'), line.end());
                line.erase(std::remove(line.begin(), line.end(), '\r'), line.end());
                alphabet.push_back(line);
            }
        }
        alpha_size = alphabet.size();
        return true;
    }

    virtual std::vector <std::string> strDecode(std::vector<int> &preds, bool raw) {
        std::vector <std::string> str;
        if (raw) {
            for (auto v: preds) {
                str.push_back(alphabet[v]);
            }
        } else {
            for (size_t i = 0; i < preds.size(); i++) {
                if (preds[i] == 0 || (i > 0 && preds[i - 1] == preds[i])) continue;
                str.push_back(alphabet[preds[i] - 1]);
            }
        }
        return str;
    };

    virtual std::vector <std::vector<std::string>> Output(std::vector <std::vector<float>> &output, bool raw = false) {
        std::vector <std::vector<std::string>> chars(output.size());
        for (size_t batch_id = 0; batch_id < 1; batch_id++) {
            std::vector<int> preds;
            for (int i = 0; i < alpha_len; i++) {
                float *offset = output[batch_id].data() + alpha_size * i;
                std::vector<float> alpha_(offset, offset + alpha_size);
                std::vector<float>::iterator max_iterator = std::max_element(alpha_.begin(), alpha_.end());
                int maxj = std::distance(alpha_.begin(), max_iterator);
                preds.push_back(maxj);
            }
            std::vector <std::string> aa = strDecode(preds, raw);

            for (auto char_ : aa)
                std::cout << "raw: " << char_ << std::endl;

            chars.push_back(aa);
        }
        return chars;
    };

protected:
//    int channels = 3;
    int device_id = 0;
    bool trt_serialize;
    size_t alpha_size = 6625, alpha_len = 70;
    int feats_dim = alpha_len * alpha_size;
    std::vector <std::string> alphabet;
};

class CrnnRegistry{
public:

    typedef std::shared_ptr<BaseCrnn> (*Creator)(int device_id, bool trt_serialize);

    typedef std::map<std::string, Creator> CreatorRegistry;

    static CreatorRegistry& Registry() {
        static CreatorRegistry* g_registry_ = new CreatorRegistry();
        return *g_registry_;
    }

    static std::shared_ptr<BaseCrnn > CreateRecognizer(const string type, int device_id, bool trt_serialize) {
        std::cout << "Creating class " << type << std::endl;
        CreatorRegistry& registry = Registry();
        if (registry.count(type)==1)
            std::cout << "class type " << type << " not registered." << std::endl;
        return registry[type](device_id, trt_serialize);
    }

    static void AddCreator(const string type, Creator creator) {
        CreatorRegistry& registry = Registry();
        if (registry.count(type)==0)
            std::cout << "class type " << type << " already registered.";
        registry[type] = creator;
    }

private:
    CrnnRegistry(){};
    std::mutex register_mutex_;
};

class CrnnRegisterer {
public:
    CrnnRegisterer(const string& type,
                    std::shared_ptr<BaseCrnn> (*creator)(int device_id, bool trt_serialize)) {
        std::cout << "Registering class type: " << type << std::endl;
        CrnnRegistry::AddCreator(type, creator);
    }
};

#define REGISTER_CRNN_CREATOR(type, creator)                                         \
  static CrnnRegisterer g_creator_f_##type(#type, creator);                          \

#define REGISTER_CRNN_CLASS(type)                                                     \
  std::shared_ptr<BaseCrnn> Creator_##type##Class(int device_id, bool trt_serialize)  \
  {                                                                                   \
    return std::shared_ptr<BaseCrnn>(new type(device_id, trt_serialize));             \
  }                                                                                   \
  REGISTER_CRNN_CREATOR(type, Creator_##type##Class);


#endif //OCR_BASE_CRNN_H
